#include <thrust/copy.h>
#include <cusp/precond/bspai.h>
#include <cusp/krylov/gmres.h>
#include <cusp/csr_matrix.h>
#include <cusp/io/matrix_market.h>
#include <iostream>
#include <cusp/frobenius_norm.h>
#include <cusp/multiply.h>
#include <cusp/gallery/eye.h>
#include <cusp/elementwise.h>
#include <assert.h>

template <typename Monitor>
void report_status(Monitor& monitor)
{
    if (monitor.converged())
    {
        std::cout << "Solver converged to " << monitor.tolerance() << " tolerance";
        std::cout << " after " << monitor.iteration_count() << " iterations";
        std::cout << " (" << monitor.residual_norm() << " final residual)" << std::endl;
    }
    else
    {
        std::cout << "Solver reached iteration limit " << monitor.iteration_limit() << " before converging";
        std::cout << " to " << monitor.tolerance() << " tolerance ";
        std::cout << " (" << monitor.residual_norm() << " final residual)" << std::endl;
    }
}

int main(int argc, char** argv)
{
    typedef unsigned int        IndexType;
    typedef unsigned int        IndexType2;
    typedef float               ValueType;
    typedef cusp::device_memory MemorySpace;


    std::cout << "Loading..." << std::endl;
    // create an empty sparse matrix structure
    cusp::csr_matrix<IndexType, ValueType, cusp::host_memory> A;
    cusp::array2d<ValueType, MemorySpace> rhs(A.num_rows, 1);
    cusp::csr_matrix<IndexType, IndexType2, cusp::host_memory> A_BSP;
    
    cusp::io::read_matrix_market_file(A, std::string(argv[1]));
    if(argc>3)
        cusp::io::read_matrix_market_file(rhs, std::string(argv[3]));
    cusp::io::read_matrix_market_file(A_BSP, std::string(argv[2]));
    
    std::cout << "Matrix size: " << A.num_rows << " with " << A.num_entries << " non-zeros." << std::endl;
    std::cout << "Block sparsity pattern size: " << A_BSP.num_rows << " with " << A_BSP.num_entries << " non-zero blocks." << std::endl;

    cusp::csr_matrix<IndexType, ValueType, MemorySpace> d_A(A);
    cusp::csr_matrix<IndexType, IndexType2, MemorySpace> d_A_BSP(A_BSP);
    
    std::cout << "Finished loading" << std::endl;

    std::cout << "Starting preconditioner" << std::endl;
    cusp::precond::bspai<IndexType, ValueType, MemorySpace> d_M(d_A, d_A_BSP);
   

    
    d_M.matrix.sort_by_row(); // XXX
    std::cout << "Completed preconditioner" << std::endl;
    
    /*
   	cusp::csr_matrix<IndexType, ValueType, MemorySpace> d_I;
    cusp::gallery::eye(d_I, A.num_rows, A.num_cols);
    cusp::csr_matrix<IndexType,ValueType, MemorySpace> d_P;
    cusp::multiply(d_M.matrix, d_A, d_P);
    cusp::subtract(d_P, d_I, d_P);
    std::cout << "#Fnorm(A)\t" << cusp::frobenius_norm(d_A) << "\n";
    std::cout << "#Fnorm(MA-I)\t" << cusp::frobenius_norm(d_P) << std::endl;
    */
    
    std::string newFname;
    std::string fname = std::string(argv[2]);
    std::string ext(".mtx");
    
    //newFname=fname.substr(0, fname.find(ext)).append("_BSPAI.mtx");
    //cusp::coo_matrix<IndexType, ValueType, cusp::host_memory> M(d_M.matrix);
    //cusp::io::write_matrix_market_file(M, newFname);
        
    
    int restart=1500;
    
    {
        std::cout << "\nSolving with no preconditioner" << std::endl;
    
        // allocate storage for solution (x) and right hand side (b)
        cusp::array1d<ValueType, MemorySpace> d_x(d_A.num_rows, 0);
        cusp::array1d<ValueType, MemorySpace> d_b(d_A.num_rows, 1);
        thrust::copy(rhs.values.begin(),rhs.values.end(),d_b.begin());
        
        // cusp::print(b);
        // set stopping criteria (iteration_limit = 1000, relative_tolerance = 1e-6)
        cusp::verbose_monitor<ValueType> monitor(d_b, 1000, 1e-6);
        int restart=1500;
        // solve
        cusp::krylov::gmres(d_A, d_x, d_b, restart, monitor);

        // report status
        report_status(monitor);
    }
    
    /*
    {
        std::cout << "\nSolving with BSPAI preconditioner" << std::endl;
    
        // allocate storage for solution (x) and right hand side (b)
        cusp::array1d<ValueType, MemorySpace> d_x(d_A.num_rows, 0);
        cusp::array1d<ValueType, MemorySpace> d_b(d_A.num_rows, 1);
        thrust::copy(rhs.values.begin(),rhs.values.end(),d_b.begin());
        
        // cusp::print(b);
        // set stopping criteria (iteration_limit = 1000, relative_tolerance = 1e-6)
        cusp::verbose_monitor<ValueType> monitor(d_b, 1500, 1e-6);
        int restart=1500;
        // solve
        cusp::krylov::gmres(d_A, d_x, d_b, restart, monitor, d_M);

        // report status
        report_status(monitor);
    }
    */
    
    return 0;
}



